#!/usr/bin/env python3
import socket
import sys
from scapy.all import (
    IP,
    UDP,
    Ether,
    Packet,
    get_if_hwaddr,
    get_if_list,
    sendp,
    bind_layers
)
from scapy.fields import BitField, XShortField

def get_if():
    ifs = get_if_list()
    for i in ifs:
        if "eth0" in i:
            return i
    print("Cannot find eth0 interface")
    exit(1)

class SourceRoute(Packet):
    fields_desc = [BitField("bos", 0, 1),
                   BitField("port", 0, 15)]

# 绑定协议栈
bind_layers(Ether, SourceRoute, type=0x1234)
bind_layers(SourceRoute, SourceRoute, bos=0)
bind_layers(SourceRoute, IP, bos=1)

def send_ip_packet(iface, dst_ip):
    pkt = Ether(src=get_if_hwaddr(iface), dst='ff:ff:ff:ff:ff:ff') / \
          IP(dst=dst_ip) / UDP(dport=4321, sport=1234)
    print("Sending normal IPv4 packet")
    pkt.show2()
    sendp(pkt, iface=iface, verbose=False)

def send_sr_packet(iface, dst_ip, ports):
    pkt = Ether(src=get_if_hwaddr(iface), dst='ff:ff:ff:ff:ff:ff')
    n = len(ports)
    for i, p in enumerate(ports):
        bos = 1 if i == n - 1 else 0
        pkt = pkt / SourceRoute(bos=bos, port=int(p))
    pkt = pkt / IP(dst=dst_ip) / UDP(dport=4321, sport=1234)
    print("Sending source-routed packet")
    pkt.show2()
    sendp(pkt, iface=iface, verbose=False)

def main():
    if len(sys.argv) < 2:
        print('Usage: send.py <destination_ip>')
        exit(1)
    addr = socket.gethostbyname(sys.argv[1])
    iface = get_if()
    print("Sending on interface %s to %s" % (iface, addr))

    while True:
        print("\nChoose mode:")
        print("  ip  : send normal IPv4 packet")
        print("  sr  : send source-routed packet (enter space-separated ports)")
        print("  q   : quit")
        choice = input("Enter choice: ").strip().lower()
        if choice == "q":
            break
        elif choice == "ip":
            send_ip_packet(iface, addr)
        elif choice == "sr":
            s = input('Type space-separated port nums (e.g., "2 3 2 1"): ')
            ports = s.split()
            if not ports or ports == ['']:
                print("No ports given.")
                continue
            send_sr_packet(iface, addr, ports)
        else:
            print("Invalid choice.")

if __name__ == '__main__':
    main()
